#![feature(iter_next_chunk)]
use anyhow::Result;
use async_process::{Command, Stdio};
use async_std::{print, println};
use clap::Parser;
use futures::future::join_all;
use std::time::Instant;

#[derive(Parser, Debug)]
struct ShellArgs {
    prog: String,
    args: Vec<String>,
}

#[derive(clap::ValueEnum, Clone, Debug)]
enum Shell {
    #[cfg(target_os = "windows")]
    Cmd,
    #[cfg(target_os = "windows")]
    Pwsh,
    #[cfg(not(target_os = "windows"))]
    Sh,
    Nu,
}

impl Shell {
    pub fn to_shell(&self) -> ShellArgs {
        match self {
            #[cfg(target_os = "windows")]
            Shell::Cmd => ShellArgs {
                prog: "cmd".to_string(),
                args: vec!["-c".to_string()],
            },
            #[cfg(target_os = "windows")]
            Shell::Pwsh => ShellArgs {
                prog: "pwsh".to_string(),
                args: vec!["-c".to_string()],
            },
            #[cfg(not(target_os = "windows"))]
            Shell::Sh => ShellArgs {
                prog: "sh".to_string(),
                args: vec!["-c".to_string()],
            },
            Shell::Nu => ShellArgs {
                prog: "nu".to_string(),
                args: vec!["-c".to_string()],
            },
        }
    }
}

#[derive(Debug, Parser)]
#[clap(name = "timeit")]
#[clap(bin_name = "timeit")]
#[clap(author, version, about="测试程序", long_about = None)]
struct App {
    #[arg(short, long, default_value_t = 1000, value_parser=clap::value_parser!(u128), help = "命令执行次数")]
    times: u128,
    #[arg(short, long, default_value_t = 5, value_parser=clap::value_parser!(u128), help = "重复次数")]
    repeat: u128,
    #[arg(short, long, default_value_t = 10, value_parser=clap::value_parser!(u128), help = "分块执行次数")]
    chunk: u128,
    #[cfg(target_os = "windows")]
    #[arg(short, long, default_value_t = Shell::Cmd, value_enum, help = "shell")]
    shell: Shell,
    #[cfg(not(target_os = "windows"))]
    #[arg(short, long, default_value_t = Shell::Sh, value_enum, help = "shell")]
    shell: Shell,
    commands: Vec<String>,
}

async fn run_command(shell: Shell, cmd: Vec<String>) -> Result<u128> {
    let start = Instant::now();
    let shell_args = &shell.to_shell();
    Command::new(&shell_args.prog)
        .args(&shell_args.args)
        .args(&cmd)
        .stdin(Stdio::null())
        .stdout(Stdio::null())
        .stderr(Stdio::null())
        .output()
        .await?;
    Ok(start.elapsed().as_millis())
}

#[async_std::main]
async fn main() {
    let args: App = App::parse();
    let commands = args.commands;
    let times = args.times;
    for r in 0..args.repeat {
        let mut results: Vec<Result<u128>> = vec![];
        print!("第{}轮: ", r + 1).await;
        let mut r = 0..times;
        loop {
            let chunk = r.next_chunk::<10>();
            if chunk.is_err() {
                break;
            }
            let mut chunk_futures = Vec::new();
            for _ in chunk.unwrap() {
                let c = run_command(args.shell.clone(), commands.to_owned());
                chunk_futures.push(c);
            }
            let mut chunk_results = join_all(chunk_futures).await;
            results.append(&mut chunk_results);
        }
        let ok_res = results
            .into_iter()
            .filter(|r| r.is_ok())
            .map(|r| r.unwrap())
            .collect::<Vec<u128>>();

        let total_time = ok_res.iter().sum::<u128>();
        let max_time = ok_res.iter().max().unwrap_or(&0u128);
        let min_time = ok_res.iter().min().unwrap_or(&0u128);
        println!(
            "总耗时：{}ms, 平均执行时间: {} ms, 最好：{}ms, 最差: {}ms",
            total_time,
            total_time / times,
            min_time,
            max_time
        )
        .await;
    }
}
